# Copyright 2022 The XFL Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from transformers import Trainer
from service.fed_config import FedConfig
from algorithm.framework.horizontal.chatglm.common import Common
from algorithm.framework.horizontal.chatglm.callback import LabelTrainerCallback


class HorizontalChatglmLabelTrainer(Common):
    def __init__(self, train_conf: dict):
        super().__init__(train_conf)

    def fit(self):
        agg_steps = self.common_config.aggregation.get("agg_steps") or self.common_config.train_params["trainer"]["save_steps"]
        sec_conf = self.common_config.train_params["encryption"]
        (peft_type, peft_config_dict), = self.common_config.train_params["peft"].items()
        
        my_callback = LabelTrainerCallback(agg_steps,
                                           sec_conf,
                                           root_id=FedConfig.get_assist_trainer(),
                                           leaf_ids=FedConfig.get_label_trainer(),
                                           init_params=not self.load_from_pretrained,
                                           peft_type=peft_type)
        
        trainer = Trainer(
            model=self.model,
            args=self.training_args,
            train_dataset=self.train_dataset,
            eval_dataset=self.val_dataset,
            tokenizer=self.tokenizer,
            data_collator=self.data_collator,
            callbacks=[my_callback],
        )
        
        trainer.train()
